import tensorflow as tf

from src.models.helpers.get_state import (
    get_state_erroneous_channel_state_information,
    get_state_aods,
)


class ConfigSACLearner:
    """
    配置SAC（Soft Actor-Critic）学习器的类。
    """

    def __init__(
            self,
            size_state,
            num_actions,
    ) -> None:
        """
        初始化 ConfigSACLearner 实例，并设置各项配置。

        参数:
            size_state (int): 状态空间的维度
            num_actions (int): 动作空间的维度
        """
        self.training_name: str = 'test'  # 训练名称

        # 状态获取函数及其参数
        self.get_state = get_state_erroneous_channel_state_information
        self.get_state_args = {
            'csi_format': 'rad_phase',
            'norm_csi': True,
        }

        # 经验回放缓冲区中 MMSE 动作的添加概率
        self.percentage_mmse_samples_added_to_exp_buffer: float = 0.0  # [0.0, 1.0] 的 MMSE 动作添加机会

        # 训练参数
        self.training_args: dict = {
            'future_reward_discount_gamma': 0.0,  # 未来奖励的折扣因子，用于稳定性
            'entropy_scale_alpha_initial': 1,  # '软' 熵惩罚的权重
            'target_entropy': 1.0,  # SAC 启发式实现 = 动作空间维度的乘积
            'entropy_scale_optimizer': tf.keras.optimizers.SGD,  # 熵规模优化器
            'entropy_scale_optimizer_args': {
                'learning_rate': 1e-4,  # 学习率 = 0.0 表示没有自适应熵规模，需要手动调整初始熵规模
            },
            'training_minimum_experiences': 1_000,  # 训练所需的最小经验数
            'training_batch_size': 512,  # 训练批次大小
            'training_target_update_momentum_tau': 0,  # 主网络副本更新到目标网络的比重
        }
        
        # 经验回放缓冲区参数
        self.experience_buffer_args: dict = {
            'buffer_size': 10_000,  # 缓冲区大小
            'priority_scale_alpha': 0.0,  # [0, 1] 中的 alpha，alpha=0 表示均匀采样，1 表示完全优先级采样
            'importance_sampling_correction_beta': 1.0  # [0%, 100%] 中的 beta，beta=100% 表示完全修正
        }
        
        # 网络参数
        self.network_args: dict = {
            'value_network_args': {
                'hidden_layer_units': [512, 512, 512, 512],  # 值网络的隐藏层单元数
                'activation_hidden': 'tanh',  # 激活函数，'relu'、'tanh' 或 'penalized_tanh'
                'kernel_initializer_hidden': 'glorot_uniform'  # 内核初始化器，'glorot_uniform' 或 'he_uniform'
            },
            'value_network_optimizer': tf.keras.optimizers.Adam,  # 值网络优化器
            'value_network_optimizer_args': {
                'learning_rate': 1e-5,  # 学习率
                'amsgrad': False,  # 是否使用 AMSGrad
            },
            'policy_network_args': {
                'hidden_layer_units': [512, 512, 512, 512],  # 策略网络的隐藏层单元数
                'activation_hidden': 'tanh',  # 激活函数，'relu'、'tanh' 或 'penalized_tanh'
                'kernel_initializer_hidden': 'glorot_uniform'  # 内核初始化器，'glorot_uniform' 或 'he_uniform'
            },
            'policy_network_optimizer': tf.keras.optimizers.Adam,  # 策略网络优化器
            'policy_network_optimizer_args': {
                'learning_rate': 1e-6,  # 学习率
                'amsgrad': True,  # 是否使用 AMSGrad
            },
        }

        # 训练设置
        self.training_episodes: int = 3_000  # 训练集的总回合数，每回合是一次完整的环境重置
        self.training_steps_per_episode: int = 1_000  # 每回合的训练步数

        # 每多少步训练一次策略
        self.train_policy_every_k_steps: int = 1  # 每 k 步训练一次策略，以使值函数逼近有时间稳定下来
        # 训练策略的延迟步数
        self.train_policy_after_j_steps: int = 0  # 值函数开始合理后才开始训练策略

        # 在初始化后进一步配置
        self._post_init(num_actions=num_actions, size_state=size_state)

    def _post_init(
            self,
            num_actions,
            size_state,
    ) -> None:
        """
        初始化后进一步设置网络参数和算法参数。

        参数:
            num_actions (int): 动作空间的维度
            size_state (int): 状态空间的维度
        """
        # 确保最小经验数不小于批次大小
        self.training_args['training_minimum_experiences'] = max(self.training_args['training_minimum_experiences'],
                                                                 self.training_args['training_batch_size'])
        # 设置网络参数的状态空间和动作空间维度
        self.network_args['size_state'] = size_state
        self.network_args['num_actions'] = num_actions

        # 收集的算法参数
        self.algorithm_args = {
            **self.training_args,
            'network_args': self.network_args,
            'experience_buffer_args': self.experience_buffer_args,
        }
